#include "device/vector_path.h"
#include <memory>
#define WidthSize 1024
#define BufferSize WidthSize * WidthSize

VectorPath::VectorPath(std::shared_ptr<Dram> dramobj) {
  DramPtr = std::move(dramobj);
  pe_start = false;
  pe_done = false;
  ifmBuffer[0] = new s8[BufferSize];
  ifmBuffer[1] = new s8[BufferSize];
  ofmBuffer = new s8[BufferSize];
}

VectorPath::~VectorPath() {
  delete[] ifmBuffer[0];
  delete[] ifmBuffer[1];
  delete[] ofmBuffer;
}

void VectorPath::element_wise(u32 length) {
  CHECK(length <= BufferSize);
  for (u32 i = 0; i < length; i++) {
    s64 result = 0;
    if (pe_mode == 0)
      result = ((s64)ifmBuffer[0][i] >> element_q1) +
               ((s64)ifmBuffer[1][i] >> element_q2);
    else if (pe_mode == 1)
      result = ((s64)ifmBuffer[0][i] >> element_q1) *
               ((s64)ifmBuffer[1][i] >> element_q2);
    else
      CHECK(false);
    post_value(&result, true);
    ofmBuffer[i] = (s8)result;
  }
}

void VectorPath::pooling(u64 ofm_addr, u64 ifm_addr) {
  u32 pool_w = 0, pool_h = 0, pooling_type = 0;
  switch (pe_mode) {
  case 1 << 2: // ofm = 2x2avepooling(ifm0)
    pool_w = 2;
    pool_h = 2;
    pooling_type = 0;
    break;
  case 1 << 3: // ofm = 3x3avepooling(ifm0)
    pool_w = 3;
    pool_h = 3;
    pooling_type = 0;
    break;
  case 1 << 4: // ofm = 4x4avepooling(ifm0)
    pool_w = 4;
    pool_h = 4;
    pooling_type = 0;
    break;
  case 1 << 5: // ofm = 2x2maxpooling(ifm0)
    pool_w = 2;
    pool_h = 2;
    pooling_type = 1;
    break;
  case 1 << 6: // ofm = 3x3maxpooling(ifm0)
    pool_w = 3;
    pool_h = 3;
    pooling_type = 1;
    break;
  case 1 << 7: // ofm = 4x4maxpooling(ifm0)
    pool_w = 4;
    pool_h = 4;
    pooling_type = 1;
    break;
  default:
    CHECK(false);
  }
  u32 stride_w = pooling_stride & 0xf;
  u32 stride_h = (pooling_stride & 0xf0) >> 4;
  u32 pad_left = (padding_idc & 1) ? (padding_size & 0xf) : 0;
  u32 pad_right = (padding_idc & 2) ? ((padding_size & 0xf0) >> 4) : 0;
  u32 pad_top = (padding_idc & 4) ? ((padding_size & 0xf00) >> 8) : 0;
  u32 pad_bottom = (padding_idc & 8) ? ((padding_size & 0xf000) >> 12) : 0;
  CHECK((ifm_width + pad_left + pad_right - pool_w) / stride_w + 1 == ofm_width);
  CHECK((ifm_height + pad_top + pad_bottom - pool_h) / stride_h + 1 == ofm_height);
  CHECK(ifm_num == ofm_num);
  u32 ifm_w_sum = pad_left + ifm_width + pad_right;
  u32 ifm_h_sum = pad_top + ifm_height + pad_bottom;
  //make sure ifm_w/ifm_h < WidthSize
  CHECK(ifm_w_sum < WidthSize);
  CHECK(ifm_h_sum < WidthSize);
  u64 input_addr = ifm_addr;
  u64 output_addr = ofm_addr;
  for(u32 ifm_num_index = 0;ifm_num_index < ifm_num;ifm_num_index++){
    u64 input_addr_tmp = input_addr;
    u32 index_j = 0;
    for(;index_j < pad_top; index_j++){
        for(u32 index_i =0;index_i < ifm_w_sum;index_i++){
            ifmBuffer[0][index_j * ifm_w_sum + index_i] = padding_word;
        }
    }
    for(;index_j < pad_top+ifm_height;index_j++){
        u32 index_i = 0;
        for(;index_i<pad_left;index_i++){
            ifmBuffer[0][index_j * ifm_w_sum + index_i] = padding_word;
        }
        DramPtr->read_mem(&(ifmBuffer[0][index_j * ifm_w_sum + index_i]),
                          input_addr_tmp, ifm_width);
        input_addr_tmp += ifm_width;
        index_i += ifm_width;
        for(;index_i<ifm_w_sum;index_i++){
            ifmBuffer[0][index_j * ifm_w_sum + index_i] = padding_word;
        }
    }
    for(;index_j < ifm_h_sum; index_j++){
        for(u32 index_i =0;index_i < ifm_w_sum;index_i++){
            ifmBuffer[0][index_j * ifm_w_sum + index_i] = padding_word;
        }
    }
    pooling_tile(ifm_w_sum, ifm_h_sum, ofm_width, ofm_height,
                 pool_w, pool_h, pooling_type, stride_w, stride_h);

    DramPtr->write_mem(output_addr, ofmBuffer, ofm_width * ofm_height);

    input_addr += ifm_offset[0];
    output_addr += ofm_offset;
  }
  //u32 tile_max_output_w = (WidthSize - pool_w) / stride_w + 1;
  //u32 tile_max_output_h = (WidthSize - pool_h) / stride_h + 1;
  //u32 tile_max_intput_w = (tile_max_output_w - 1) * stride_w + pool_w;
  //u32 tile_max_intput_h = (tile_max_output_h - 1) * stride_h + pool_h;
  //u32 ifm_w_sum = pad_left + ifm_width + pad_right;
  //u32 ifm_h_sum = pad_top + ifm_height + pad_bottom;
  //u32 tile_index_w = (ofm_width + tile_max_output_w - 1)/tile_max_output_w + 1;
  //u32 tile_index_h = (ofm_height + tile_max_output_h - 1)/tile_max_output_h + 1;
  //for(int index_j=0;index_j<tile_index_h - 1;index_j++){
  //    for(int index_i=0;index_i<tile_index_w - 1;index_i++){
  //        u32 start_output_i = index_i * tile_index_w;
  //        u32 end_output_i = (index_i+1) * tile_index_w;
  //        u32 start_output_j = index_j * tile_index_h;
  //        u32 end_output_j = (index_j+1) * tile_index_h;
  //        u32 start_input_i = start_output_i * stride_w;
  //        u32 end_input_i = (end_output_i - 1) * stride_w + pool_w;
  //        u32 start_input_j = start_output_j * stride_h;
  //        u32 end_input_j = (end_output_j - 1) * stride_h + pool_h;
  //    }
  //}

  return;
}

void VectorPath::pooling_tile(u32 input_w, u32 input_h, u32 output_w,
                              u32 output_h, u32 pool_w, u32 pool_h,
                              u32 pooling_type, u32 stride_w, u32 stride_h) {
  CHECK((input_w - pool_w) / stride_w + 1 == output_w);
  CHECK((input_h - pool_h) / stride_h + 1 == output_h);

  for (u32 out_j = 0; out_j < output_h; out_j++) {
    for (u32 out_i = 0; out_i < output_w; out_i++) {
      s64 result;
      if (pooling_type == 0)
        result = 0;
      else if (pooling_type == 1)
        result = -128;
      else
        CHECK(false);

      for (u32 win_j = 0; win_j < pool_h; win_j++) {
        for (u32 win_i = 0; win_i < pool_w; win_i++) {
          u32 index_w = out_i * stride_w + win_i;
          u32 index_h = out_j * stride_h + win_j;
          s8 input_value = ifmBuffer[0][index_h * input_w + index_w];
          if (pooling_type == 0)
            result += input_value;
          else {
            if (result < input_value)
              result = input_value;
          }
        }
      }
      //if (pooling_type == 0) {
      //  if (pool_w == 2 && pool_h == 2)
      //    result = result >> 2;
      //  else if (pool_w == 4 && pool_h == 4)
      //    result = result >> 4;
      //  else if (pool_w == 3 && pool_h == 3)
      //    post_value(&result, true);
      //  else
      //    CHECK(false);
      //}
      if (pooling_type == 0) {
        result = (s32)(((s64)result * re_mul_num) >> (out_shift - 1));
        if(out_shift > 0)result = (result + 1) >> 1;
        if(result > 127) result = 127;
        if(result < -127) result = -127;
      }
 
      ofmBuffer[out_j * output_w + out_i] = (s8)result;
    }
  }
}

void VectorPath::reduce_sum(u32 length, s64 *result) {
  CHECK(length <= BufferSize);
  for (u32 i = 0; i < length; i++) {
    *result += ifmBuffer[0][i];
  }
  return;
}

void VectorPath::arg_max(u32 length, u32 index_base, u32 *max_index,
                         s8 *max_value) {
  CHECK(length <= BufferSize);
  for (u32 i = 0; i < length; i++) {
    if (*max_value < ifmBuffer[0][i]) {
      *max_value = ifmBuffer[0][i];
      *max_index = index_base + i;
    }
  }
  return;
}

void VectorPath::post_value(s64 *result, bool is_saturate) {
  *result = *result * re_mul_num;
  if (out_shift < 0)
    *result = *result << (-out_shift);
  if (out_shift > 0)
    *result = *result >> out_shift;
  *result = *result >> 31;
  *result = (*result + 1) >> 1;
  if (is_saturate) {
    if (*result > 127)
      *result = 127;
    if (*result < -127)
      *result = -127;
  }
  return;
}

void VectorPath::Run() {
  if (!pe_start)
    return;
  if (pe_done)
    CHECK(false);
  u64 ifm_addr[2];
  u64 ofm_addr;
  ifm_addr[0] = ((u64)ifm_start_addr_h[0] << 32) + ifm_start_addr_l[0];
  ifm_addr[1] = ((u64)ifm_start_addr_h[1] << 32) + ifm_start_addr_l[1];
  ofm_addr = ((u64)ofm_start_addr_h << 32) + ofm_start_addr_l;
  switch (pe_mode) {
  case 0:      // ofm = ifm0 + ifm1
  case 1 << 1: // ofm = ifm0 * ifm1
  {
    CHECK(ifm_size == ofm_size);
    CHECK(ifm_num == ofm_num);
    for (u32 map_index = 0; map_index < ifm_num; map_index++) {
      u64 ifm_addr_tmp[2];
      ifm_addr_tmp[0] = ifm_addr[0];
      ifm_addr_tmp[1] = ifm_addr[1];
      u64 ofm_addr_tmp = ofm_addr;
      u32 ifm_size_remain = ifm_size;
      while (ifm_size_remain > BufferSize) {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp[0], BufferSize);
        DramPtr->read_mem(ifmBuffer[1], ifm_addr_tmp[1], BufferSize);
        ifm_addr_tmp[0] += BufferSize;
        ifm_addr_tmp[1] += BufferSize;
        element_wise(BufferSize);
        DramPtr->write_mem(ofm_addr_tmp, ofmBuffer, BufferSize);
        ofm_addr_tmp += BufferSize;
        ifm_size_remain -= BufferSize;
      }
      {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp[0], ifm_size_remain);
        DramPtr->read_mem(ifmBuffer[1], ifm_addr_tmp[1], ifm_size_remain);
        element_wise(ifm_size_remain);
        DramPtr->write_mem(ofm_addr_tmp, ofmBuffer, ifm_size_remain);
      }

      ifm_addr[0] += ifm_offset[0];
      ifm_addr[1] += ifm_offset[1];
    }
    break;
  }
  case 1 << 2: // ofm = 2x2avepooling(ifm0)
  case 1 << 3: // ofm = 3x3avepooling(ifm0)
  case 1 << 4: // ofm = 4x4avepooling(ifm0)
  case 1 << 5: // ofm = 2x2maxpooling(ifm0)
  case 1 << 6: // ofm = 3x3maxpooling(ifm0)
  case 1 << 7: // ofm = 4x4maxpooling(ifm0)
  {
    pooling(ofm_addr, ifm_addr[0]);
    break;
  }
  case 1 << 8: // ofm = reduce(ifm0)
  {
    s64 result = 0;
    for (u32 map_index = 0; map_index < ifm_num; map_index++) {
      u64 ifm_addr_tmp = ifm_addr[0];
      u32 ifm_size_remain = ifm_size;
      while (ifm_size_remain > BufferSize) {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp, BufferSize);
        ifm_addr_tmp += BufferSize;
        reduce_sum(BufferSize, &result);
        ifm_size_remain -= BufferSize;
      }
      {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp, ifm_size_remain);
        reduce_sum(ifm_size_remain, &result);
      }
      ifm_addr[0] += ifm_offset[0];
    }
    post_value(&result, false);
    s32 result_s32 = (s32)result;
    DramPtr->write_mem(ofm_addr, &result_s32, sizeof(s32));
    break;
  }
  case 1 << 9: // ofm = arg_max(ifm0)
  {
    u32 index_base = 0;
    u32 max_index = -1;
    s8 max_value = -128;
    for (u32 map_index = 0; map_index < ifm_num; map_index++) {
      u64 ifm_addr_tmp = ifm_addr[0];
      u32 ifm_size_remain = ifm_size;
      while (ifm_size_remain > BufferSize) {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp, BufferSize);
        ifm_addr_tmp += BufferSize;
        arg_max(BufferSize, index_base, &max_index, &max_value);
        ifm_size_remain -= BufferSize;
        index_base += BufferSize;
      }
      {
        DramPtr->read_mem(ifmBuffer[0], ifm_addr_tmp, ifm_size_remain);
        arg_max(ifm_size_remain, index_base, &max_index, &max_value);
        index_base += ifm_size_remain;
      }
      ifm_addr[0] += ifm_offset[0];
    }
    DramPtr->write_mem(ofm_addr, &max_index, sizeof(u32));
    break;
  }
  default:
    CHECK(false);
  }
  pe_done = true;
  pe_start = false;
  return;
}
